package com.shujia.alink;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.evaluation.EvalBinaryClassBatchOp;
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.BinaryClassMetrics;
import com.alibaba.alink.pipeline.Pipeline;
import com.alibaba.alink.pipeline.PipelineModel;
import com.alibaba.alink.pipeline.classification.NaiveBayesTextClassifier;
import com.alibaba.alink.pipeline.nlp.DocCountVectorizer;
import com.alibaba.alink.pipeline.nlp.Segment;
import com.alibaba.alink.pipeline.nlp.StopWordsRemover;

public class NaiveBayes {
    public static void main(String[] args) throws Exception {
        String SCHEMA_STR = "label double, text string";

        //数据源 source
        BatchOperator data = new CsvSourceBatchOp()
                .setFilePath("flink_alink/data/train.txt")
                .setSchemaStr(SCHEMA_STR)
                .setFieldDelimiter("\t");

        /**
         * Pipeline  机器学习流程
         * Segment 中文分词（内部使用的是jieba分词器）
         *StopWordsRemover  取出停留词
         *DocCountVectorizer 转换成词典向量 加上tf-idf
         * NaiveBayesTextClassifier 分也是分类算法
         */

        Pipeline pipeline = new Pipeline()
                .add(new Segment().setSelectedCol("text"))
                .add(new StopWordsRemover().setSelectedCol("text"))
                .add(new DocCountVectorizer().setSelectedCol("text"))
                .add(new NaiveBayesTextClassifier()
                        .setPredictionCol("prediction")
                        .setPredictionDetailCol("predictionDetail")
                        .setVectorCol("text")
                        .setLabelCol("label")
                );


        //训练模型，执行流程
        PipelineModel model = pipeline.fit(data);

        //预测
        BatchOperator<?> transform = model.transform(data);

        //模型评估
        BinaryClassMetrics metrics = new EvalBinaryClassBatchOp()
                .setLabelCol("label")
                .setPredictionDetailCol("predictionDetail")
                .linkFrom(transform)
                .collectMetrics();

        System.out.println("AUC:" + metrics.getAuc());
        System.out.println("KS:" + metrics.getKs());
        System.out.println("PRC:" + metrics.getPrc());
        System.out.println("Accuracy:" + metrics.getAccuracy());//模型准确率
        System.out.println("Macro Precision:" + metrics.getMacroPrecision());
        System.out.println("Micro Recall:" + metrics.getMicroRecall());
        System.out.println("Weighted Sensitivity:" + metrics.getWeightedSensitivity());

        //保存模型
        model.save("flink_alink/data/model");

        //启动任务
        BatchOperator.execute();
    }
}
